package org.apache.lucene.ocean.database;

import java.io.IOException;
import java.io.Serializable;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutorService;

import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.distributed.index.IndexServiceImpl.MultiValueTokenStream;
import org.apache.lucene.distributed.index.IndexServiceImpl.TokenOrderingFilter;
import org.apache.lucene.distributed.index.SearchableService.HighlightKey;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.FieldSelector;
import org.apache.lucene.document.FieldSelectorResult;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.ocean.Snapshot;
import org.apache.lucene.ocean.TransactionSystem;
import org.apache.lucene.ocean.util.NamedList;
import org.apache.lucene.search.OceanMultiThreadSearcher;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.highlight.Highlighter;
import org.apache.lucene.search.highlight.TextFragment;
import org.apache.lucene.search.highlight.TokenSources;

/**
 * Performs search using multiple threads.  Also enabled for highlighting.
 * 
 *
 */
public class MultiThreadSearchTask implements Task {
  SearchKey searchKey; // main key that defines search
  HighlightKey[] highlightKeys; // array of fields that should be highlighted
  FieldSelector fieldSelector; // field selector
  
  public MultiThreadSearchTask(SearchKey searchKey, HighlightKey[] highlightKeys, FieldSelector fieldSelector) {
    this.searchKey = searchKey;
    this.highlightKeys = highlightKeys;
    this.fieldSelector = fieldSelector;
  }
  
  public static class SetFieldSelector implements FieldSelector, Serializable {
    private Set<String> set;
    
    public SetFieldSelector(String... strings) {
      set = new HashSet<String>(strings.length);
      for (String str : strings) 
        set.add(str);
    }
    
    public FieldSelectorResult accept(String fieldName) {
      if (set.contains(fieldName)) {
        return FieldSelectorResult.LOAD;
      } else return FieldSelectorResult.NO_LOAD;
    }
  }
  
  public static class Result {
    public OceanObject[] objects;
    public float maxScore;
    public int totalHits;
    public HighlightObject[] highlights;
  }
  
  public static class HighlightObject {
    public NamedList<String,TextFragment> textFragments = new NamedList<String,TextFragment>();
  }
  
  protected HighlightObject getHighlights(int docId, Document document, IndexReader reader) throws Exception {
    HighlightObject highlightObject = new HighlightObject();
    for (HighlightKey hk : highlightKeys) {
      HighlighterImpl highlighterImpl = new HighlighterImpl(reader, hk);
      TextFragment[] textFragments = highlighterImpl.dos(docId, hk.field, document);
      for (TextFragment textFragment : textFragments) {
        highlightObject.textFragments.add(hk.field, textFragment);
      }
    }
    return highlightObject;
  }
  
  /**
   * Main method that implements the functionality of search.
   */
  public Object run(TransactionSystem transactionSystem) throws Exception {
    Snapshot snapshot = transactionSystem.getSnapshots().getLatestSnapshot();
    try {
      IndexReader reader = snapshot.getIndexReader();
      ExecutorService searchThreadPool = transactionSystem.getSearchThreadPool();
      OceanMultiThreadSearcher searcher = new OceanMultiThreadSearcher(snapshot, searchThreadPool);
      TopDocs topDocs = searcher.search(searchKey.weight, searchKey.filter, searchKey.start+searchKey.num);
      List<OceanObject> objects = new ArrayList<OceanObject>(searchKey.num);
      List<HighlightObject> highlights = new ArrayList<HighlightObject>(searchKey.num);
      for (int x=searchKey.start; x < topDocs.scoreDocs.length; x++) {
        Document document = searcher.doc(topDocs.scoreDocs[x].doc, fieldSelector);
        if (highlightKeys != null) highlights.add(getHighlights(topDocs.scoreDocs[x].doc, document, reader));
        objects.add(new OceanObject(document));
      }
      Result result = new Result();
      result.maxScore = topDocs.getMaxScore();
      result.totalHits = topDocs.totalHits;
      result.objects = (OceanObject[])objects.toArray(new OceanObject[0]);
      if (highlightKeys != null) result.highlights = (HighlightObject[])highlights.toArray(new HighlightObject[0]);
      return result;
    } finally {
      snapshot.decRef();
    }
  }
  
  /**
   * Implements highlighting functionality using the contrib highlighting code
   *
   */
  public static class HighlighterImpl {
    Highlighter highlighter;
    HighlightKey highlight;
    TokenOrderingFilter tstream;
    IndexReader reader;
    
    public HighlighterImpl(IndexReader reader, HighlightKey highlight) {
      this.reader = reader;
      this.highlight = highlight;
    }
    
    public TextFragment[] dos(int docId, String field, Document doc) throws IOException {
      String[] docTexts = doc.getValues(field);
      TextFragment[] frag;
      if (docTexts.length == 1) {
        // single-valued field
        TokenStream tstream;
        try {
          // attempt term vectors
          tstream = TokenSources.getTokenStream(reader, docId, field);
        } catch (IllegalArgumentException e) {
          // fall back to analyzer
          tstream = new TokenOrderingFilter(highlight.analyzer.tokenStream(field, new StringReader(docTexts[0])), 10);
        }
        frag = highlighter.getBestTextFragments(tstream, docTexts[0], false, highlight.numFragments);
      } else {
        // multi-valued field
        MultiValueTokenStream tstream;
        tstream = new MultiValueTokenStream(field, docTexts, highlight.analyzer, true);
        frag = highlighter.getBestTextFragments(tstream, tstream.asSingleValue(), false, highlight.numFragments);
      }
      return frag;
    }
  }
}
